// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <pollux/connectors/hive/partition_id_generator.h>

#include <pollux/connectors/hive/hive_partition_util.h>
#include <pollux/dwio/catalog/hive/file_utils.h>

using namespace kumo::pollux::dwio::catalog::hive;

namespace kumo::pollux::connector::hive {

PartitionIdGenerator::PartitionIdGenerator(
    const RowTypePtr& inputType,
    std::vector<column_index_t> partitionChannels,
    uint32_t maxPartitions,
    memory::MemoryPool* pool,
    bool partitionPathAsLowerCase)
    : partitionChannels_(std::move(partitionChannels)),
      maxPartitions_(maxPartitions),
      partitionPathAsLowerCase_(partitionPathAsLowerCase) {
  POLLUX_USER_CHECK(
      !partitionChannels_.empty(), "There must be at least one partition key.");
  for (auto channel : partitionChannels_) {
    hashers_.emplace_back(
        exec::VectorHasher::create(inputType->childAt(channel), channel));
  }

  std::vector<TypePtr> partitionKeyTypes;
  std::vector<std::string> partitionKeyNames;
  for (auto channel : partitionChannels_) {
    POLLUX_USER_CHECK(
        exec::VectorHasher::typeKindSupportsValueIds(
            inputType->childAt(channel)->kind()),
        "Unsupported partition type: {}.",
        inputType->childAt(channel)->toString());
    partitionKeyTypes.push_back(inputType->childAt(channel));
    partitionKeyNames.push_back(inputType->nameOf(channel));
  }

  partitionValues_ = BaseVector::create<RowVector>(
      ROW(std::move(partitionKeyNames), std::move(partitionKeyTypes)),
      maxPartitions_,
      pool);
  for (auto& key : partitionValues_->children()) {
    key->resize(maxPartitions_);
  }
}

void PartitionIdGenerator::run(
    const RowVectorPtr& input,
    raw_vector<uint64_t>& result) {
  const auto numRows = input->size();
  result.resize(numRows);

  // Compute value IDs using VectorHashers and store these in 'result'.
  computeValueIds(input, result);

  // Convert value IDs in 'result' into partition IDs using partitionIds
  // mapping. Update 'result' in place.

  // TODO Optimize common use case where all records belong to the same
  // partition. VectorHashers keep track of the number of unique values, hence,
  // we can find out if there is only one unique value for each partition key.
  for (auto i = 0; i < numRows; ++i) {
    auto valueId = result[i];
    auto it = partitionIds_.find(valueId);
    if (it != partitionIds_.end()) {
      result[i] = it->second;
    } else {
      uint64_t nextPartitionId = partitionIds_.size();
      POLLUX_USER_CHECK_LT(
          nextPartitionId,
          maxPartitions_,
          "Exceeded limit of {} distinct partitions.",
          maxPartitions_);

      partitionIds_.emplace(valueId, nextPartitionId);
      savePartitionValues(nextPartitionId, input, i);

      result[i] = nextPartitionId;
    }
  }
}

std::string PartitionIdGenerator::partitionName(uint64_t partitionId) const {
  return FileUtils::makePartName(
      extractPartitionKeyValues(partitionValues_, partitionId),
      partitionPathAsLowerCase_);
}

void PartitionIdGenerator::computeValueIds(
    const RowVectorPtr& input,
    raw_vector<uint64_t>& valueIds) {
  allRows_.resize(input->size());
  allRows_.setAll();

  bool rehash = false;
  for (auto& hasher : hashers_) {
    // NOTE: for boolean column type, computeValueIds() always returns true and
    // this might cause problem in case of multiple boolean partition columns as
    // we might not set the multiplier properly.
    auto partitionVector = input->childAt(hasher->channel())->loaded_vector();
    hasher->decode(*partitionVector, allRows_);
    if (!hasher->computeValueIds(allRows_, valueIds)) {
      rehash = true;
    }
  }

  if (!rehash && hasMultiplierSet_) {
    return;
  }

  uint64_t multiplier = 1;
  for (auto& hasher : hashers_) {
    hasMultiplierSet_ = true;
    multiplier = hasher->type_kind() == TypeKind::BOOLEAN
        ? hasher->enableValueRange(multiplier, 50)
        : hasher->enableValueIds(multiplier, 50);

    POLLUX_CHECK_NE(
        multiplier,
        exec::VectorHasher::kRangeTooLarge,
        "Number of requested IDs is out of range.");
  }

  for (auto& hasher : hashers_) {
    const bool ok = hasher->computeValueIds(allRows_, valueIds);
    POLLUX_CHECK(ok);
  }

  updateValueToPartitionIdMapping();
}

void PartitionIdGenerator::updateValueToPartitionIdMapping() {
  if (partitionIds_.empty()) {
    return;
  }

  const auto numPartitions = partitionIds_.size();

  partitionIds_.clear();

  raw_vector<uint64_t> newValueIds(numPartitions);
  SelectivityVector rows(numPartitions);
  for (auto i = 0; i < hashers_.size(); ++i) {
    auto& hasher = hashers_[i];
    hasher->decode(*partitionValues_->childAt(i), rows);
    const bool ok = hasher->computeValueIds(rows, newValueIds);
    POLLUX_CHECK(ok);
  }

  for (auto i = 0; i < numPartitions; ++i) {
    partitionIds_.emplace(newValueIds[i], i);
  }
}

void PartitionIdGenerator::savePartitionValues(
    uint64_t partitionId,
    const RowVectorPtr& input,
    vector_size_t row) {
  for (auto i = 0; i < partitionChannels_.size(); ++i) {
    auto channel = partitionChannels_[i];
    partitionValues_->childAt(i)->copy(
        input->childAt(channel).get(), partitionId, row, 1);
  }
}

} // namespace kumo::pollux::connector::hive
